package org.apache.lucene.search;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.locks.ReentrantLock;

import org.apache.lucene.ocean.OceanSearcher;
import org.apache.lucene.ocean.Snapshot;
import org.apache.lucene.ocean.util.Util;
import org.apache.lucene.util.PriorityQueue;

/**
 * Implements a multi threaded searcher using an ExecutorService thread pool.
 * Methods that use HitCollector may experience odd behavior because the 
 * HitCollector.collect(int doc, float score) would be called concurrently.  
 *
 */
public class OceanMultiThreadSearcher extends OceanSearcher {
  private Snapshot snapshot;
  private ExecutorService searchThreadPool;
  private Searchable[] searchables;
  private int[] starts;

  public OceanMultiThreadSearcher(Snapshot snapshot, ExecutorService searchThreadPool) throws IOException {
    super(snapshot);
    searchables = snapshot.getSearchables();
    starts = snapshot.getStarts();
    this.snapshot = snapshot;
    this.searchThreadPool = searchThreadPool;
  }

  public TopDocs search(Weight weight, Filter filter, int nDocs) throws IOException {
    HitQueue hq = new HitQueue(nDocs);
    int totalHits = 0;
    MultiSearcherThread[] msta = new MultiSearcherThread[searchables.length];
    ReentrantLock lock = new ReentrantLock();
    List<Future> futures = new ArrayList<Future>(searchables.length);
    for (int i = 0; i < searchables.length; i++) { // search each searcher
      msta[i] = new MultiSearcherThread(searchables[i], weight, filter, nDocs, hq, i, starts, lock);
      futures.add(searchThreadPool.submit(msta[i]));
    }
    processFutures(futures);
    for (int i = 0; i < searchables.length; i++) {
      IOException ioe = msta[i].getIOException();
      if (ioe == null) {
        totalHits += msta[i].hits();
      } else {
        // if one search produced an IOException, rethrow it
        throw ioe;
      }
    }
    ScoreDoc[] scoreDocs = new ScoreDoc[hq.size()];
    for (int i = hq.size() - 1; i >= 0; i--)
      // put docs in array
      scoreDocs[i] = (ScoreDoc) hq.pop();
    float maxScore = (totalHits == 0) ? Float.NEGATIVE_INFINITY : scoreDocs[0].score;
    return new TopDocs(totalHits, scoreDocs, maxScore);
  }

  private void processFutures(List<Future> futures) throws IOException {
    for (Future future : futures) {
      try {
        future.get();
      } catch (ExecutionException executionException) {
        if (executionException.getCause() instanceof IOException) {
          throw (IOException)executionException.getCause();
        }
        throw Util.asIOException(executionException.getCause());
      } catch (InterruptedException interruptedException) {
        throw Util.asIOException(interruptedException);
      }
    }
  }

  /**
   * A search implementation allowing sorting which spans a new thread for each
   * Searchable, waits for each search to complete and merges the results back
   * together.
   */
  public TopFieldDocs search(Weight weight, Filter filter, int nDocs, Sort sort) throws IOException {
    // don't specify the fields - we'll wait to do this until we get results
    FieldDocSortedHitQueue hq = new FieldDocSortedHitQueue(null, nDocs);
    int totalHits = 0;
    MultiSearcherThread[] msta = new MultiSearcherThread[searchables.length];
    ReentrantLock lock = new ReentrantLock();
    List<Future> futures = new ArrayList<Future>(searchables.length);
    for (int i = 0; i < searchables.length; i++) { // search each searcher
      // Assume not too many searchables and cost of creating a thread is by far
      // inferior to a search
      msta[i] = new MultiSearcherThread(searchables[i], weight, filter, nDocs, hq, sort, i, starts, lock);
      futures.add(searchThreadPool.submit(msta[i]));
    }
    processFutures(futures);
    float maxScore = Float.NEGATIVE_INFINITY;
    for (int i = 0; i < searchables.length; i++) {
      IOException ioe = msta[i].getIOException();
      if (ioe == null) {
        totalHits += msta[i].hits();
        maxScore = Math.max(maxScore, msta[i].getMaxScore());
      } else {
        // if one search produced an IOException, rethrow it
        throw ioe;
      }
    }
    ScoreDoc[] scoreDocs = new ScoreDoc[hq.size()];
    for (int i = hq.size() - 1; i >= 0; i--)
      // put docs in array
      scoreDocs[i] = (ScoreDoc) hq.pop();

    return new TopFieldDocs(totalHits, scoreDocs, hq.getFields(), maxScore);
  }

  class MultiSearcherThread implements Callable {
    private Searchable searchable;
    private Weight weight;
    private Filter filter;
    private int nDocs;
    private TopDocs docs;
    private int i;
    private PriorityQueue hq;
    private int[] starts;
    private IOException ioe;
    private Sort sort;
    private ReentrantLock lock;

    public MultiSearcherThread(Searchable searchable, Weight weight, Filter filter, int nDocs, HitQueue hq, int i, int[] starts, ReentrantLock lock) {
      this.searchable = searchable;
      this.weight = weight;
      this.filter = filter;
      this.nDocs = nDocs;
      this.hq = hq;
      this.i = i;
      this.starts = starts;
      this.lock = lock;
    }

    public MultiSearcherThread(Searchable searchable, Weight weight, Filter filter, int nDocs, FieldDocSortedHitQueue hq, Sort sort, int i,
        int[] starts, ReentrantLock lock) {
      this.searchable = searchable;
      this.weight = weight;
      this.filter = filter;
      this.nDocs = nDocs;
      this.hq = hq;
      this.i = i;
      this.starts = starts;
      this.sort = sort;
      this.lock = lock;
    }

    public Object call() throws Exception {
      try {
        docs = (sort == null) ? searchable.search(weight, filter, nDocs) : searchable.search(weight, filter, nDocs, sort);
      }
      // Store the IOException for later use by the caller of this thread
      catch (IOException ioe) {
        this.ioe = ioe;
      }
      if (ioe == null) {
        // if we are sorting by fields, we need to tell the field sorted hit
        // queue
        // the actual type of fields, in case the original list contained AUTO.
        // if the searchable returns null for fields, we'll have problems.
        if (sort != null) {
          ((FieldDocSortedHitQueue) hq).setFields(((TopFieldDocs) docs).fields);
        }
        ScoreDoc[] scoreDocs = docs.scoreDocs;
        for (int j = 0; j < scoreDocs.length; j++) { // merge scoreDocs into hq
          ScoreDoc scoreDoc = scoreDocs[j];
          scoreDoc.doc += starts[i]; // convert doc
          // it would be so nice if we had a thread-safe insert
          lock.lock();
          try {
            if (!hq.insert(scoreDoc))
              break;
          } finally {
            lock.unlock();// no more scores > minScore
          }
        }
      }
      return null;
    }

    public int hits() {
      return docs.totalHits;
    }

    public float getMaxScore() {
      return docs.getMaxScore();
    }

    public IOException getIOException() {
      return ioe;
    }
  }

  /**
   * Lower-level search API.
   * 
   * <p>
   * {@link HitCollector#collect(int,float)} is called for every non-zero
   * scoring document.
   * 
   * <p>
   * Applications should only use this if they need <i>all</i> of the matching
   * documents. The high-level search API ({@link Searcher#search(Query)}) is
   * usually more efficient, as it skips non-high-scoring hits.
   * 
   * @param weight
   *          to match documents
   * @param filter
   *          if non-null, a bitset used to eliminate some documents
   * @param results
   *          to receive hits
   * 
   */
  public void search(Weight weight, Filter filter, final HitCollector results) throws IOException {
    List<Future> futures = new ArrayList<Future>(searchables.length);
    for (int i = 0; i < searchables.length; i++) {
      final int start = starts[i];
      HitCollectorThread hitCollectorThread = new HitCollectorThread(start, searchables[i], weight, filter, results);
      futures.add(searchThreadPool.submit(hitCollectorThread));
      //searchables[i].search(weight, filter, new HitCollector() {
      //  public void collect(int doc, float score) {
      //    results.collect(doc + start, score);
      //  }
      //});
    }
    processFutures(futures);
  }

  private class HitCollectorThread implements Callable {
    private Searchable searchable;
    private Weight weight;
    private Filter filter;
    private HitCollector hitCollector;
    private int start;
    
    public HitCollectorThread(int start, Searchable searchable, Weight weight, Filter filter, HitCollector hitCollector) {
      this.start = start;
      this.searchable = searchable;
      this.weight = weight;
      this.filter = filter;
      this.hitCollector = hitCollector;
    }

    public Object call() throws Exception {
      searchable.search(weight, filter, new HitCollector() {
        public void collect(int doc, float score) {
          hitCollector.collect(doc + start, score);
        }
      });
      return null;
    }
  }
}
